import os 
from FlagEmbedding import FlagModel
import sys
# 添加父级目录到路径
current_dir = os.path.dirname(os.path.abspath(__file__))
parent_dir = os.path.dirname(current_dir)
sys.path.insert(0, parent_dir)
from my_common import BGE_LARGE_ZH_V1_5_MODEL


os.environ['TRANSFORMERS_NO_ADVISORY_WARNINGS'] = 'true'
# single GPU is better for small tasks
os.environ['CUDA_VISIBLE_DEVICES'] = '0'

if __name__ == '__main__':
    # get the BGE embedding model
    model = FlagModel(BGE_LARGE_ZH_V1_5_MODEL,
        query_instruction_for_retrieval="Represent this sentence for searching relevant passages:",
        query_instruction_format='{}{}',)
    
    queries = ["query 1", "query 2"]
    corpus = ["passage 1", "passage 2"]

    # encode the queries and corpus
    q_embeddings = model.encode_queries(queries)
    p_embeddings = model.encode_corpus(corpus)

    # compute the similarity scores
    scores = q_embeddings @ p_embeddings.T
    print(scores)
